#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Similarity calculation functions for collaborative filtering recommendation.

Created on: 2025-04-18
Author: Nianqing Liu
"""

import numpy as np
from sklearn.metrics.pairwise import pairwise_distances
import scipy.sparse as sp


def cosine_similarity(ratings_matrix):
    """
    Calculate cosine similarity between users.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        User-item rating matrix.

    Returns:
    --------
    numpy.ndarray
        User-user similarity matrix.
    """
    # Check if the matrix is sparse
    if sp.issparse(ratings_matrix):
        similarity = 1 - pairwise_distances(ratings_matrix, metric="cosine", n_jobs=-1)
    else:
        similarity = 1 - pairwise_distances(ratings_matrix, metric="cosine", n_jobs=-1)

    # Handle NaN values (if any)
    similarity = np.nan_to_num(similarity)

    return similarity


def pearson_similarity(ratings_matrix):
    """
    Calculate Pearson correlation between users.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        User-item rating matrix.

    Returns:
    --------
    numpy.ndarray
        User-user similarity matrix.
    """
    # Calculate correlation coefficient
    similarity = np.corrcoef(ratings_matrix)

    # Handle NaN values (if any)
    similarity = np.nan_to_num(similarity)

    return similarity


def euclidean_similarity(ratings_matrix):
    """
    Calculate similarity based on Euclidean distance between users.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        User-item rating matrix.

    Returns:
    --------
    numpy.ndarray
        User-user similarity matrix.
    """
    # Calculate Euclidean distance
    distances = pairwise_distances(ratings_matrix, metric="euclidean", n_jobs=-1)

    # Convert distance to similarity
    similarity = 1 / (1 + distances)

    return similarity


def adjusted_cosine_similarity(ratings_matrix):
    """
    Calculate adjusted cosine similarity between users.
    Normalizes ratings by subtracting user mean before calculation.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        User-item rating matrix.

    Returns:
    --------
    numpy.ndarray
        User-user similarity matrix.
    """
    # Get non-zero elements to calculate mean
    rated_mask = ratings_matrix != 0

    # Calculate user means (only for rated items)
    user_means = np.sum(ratings_matrix, axis=1) / np.sum(rated_mask, axis=1)

    # Normalize ratings by subtracting user mean (only for rated items)
    normalized_matrix = np.zeros_like(ratings_matrix)
    for i, (user_ratings, mask, mean) in enumerate(
        zip(ratings_matrix, rated_mask, user_means)
    ):
        normalized_matrix[i, mask] = user_ratings[mask] - mean

    # Calculate cosine similarity
    similarity = 1 - pairwise_distances(normalized_matrix, metric="cosine", n_jobs=-1)

    # Handle NaN values (if any)
    similarity = np.nan_to_num(similarity)

    return similarity


def calculate_similarity(ratings_matrix, method="cosine"):
    """
    Calculate user similarity matrix using the specified method.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        User-item rating matrix.
    method : str
        Similarity method to use: 'cosine', 'pearson', 'euclidean', or 'adjusted_cosine'.

    Returns:
    --------
    numpy.ndarray
        User-user similarity matrix.
    """
    if method == "cosine":
        return cosine_similarity(ratings_matrix)
    elif method == "pearson":
        return pearson_similarity(ratings_matrix)
    elif method == "euclidean":
        return euclidean_similarity(ratings_matrix)
    elif method == "adjusted_cosine":
        return adjusted_cosine_similarity(ratings_matrix)
    else:
        raise ValueError(
            f"Unknown similarity method: {method}. Choose from 'cosine', 'pearson', 'euclidean', or 'adjusted_cosine'."
        )
